#!/usr/bin/env python
# coding: utf8
"""
Celery 异步操作Ansible 服务端
"""

import os
import sys
from celery import Celery
from ansible_files.ansibleApi import *
from tools.config import BACKEND, BROKER, REDIS_ADDR, REDIS_PORT, REDIS_PD, ansible_result_redis_db, inventory
from celery.app.task import Task
from celery.utils.log import get_task_logger
from celery.result import AsyncResult
# 为了交叉导入，设置环境变量
import django
path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, path)
os.environ['DJANGO_SETTINGS_MODULE'] = 'ansible_ui.settings'
django.setup()
from public.models import AnsibleTasks, AnsiblePlaybooks

if os.environ.get("PYTHONOPTIMIZE", ""):
    print("开始启动")
else:
    print("\33[31m环境变量问题，Celery Client启动后无法正常执行Ansible任务，"
          "\n请设置export PYTHONOPTIMIZE=1；\n\33[32m程序退出\33[0m")
    sys.exit()
celery_logger = get_task_logger(__name__)

appCelery = Celery("tasks", broker=BROKER, backend=BACKEND, )
sources = inventory


# 即使 celery 执行 ansible 时有失败的任务，也要将执行的 ansible 结果保存到 db
class MyTask(Task):  # 回调
    def on_success(self, retval, task_id, args, kwargs):
        print("执行成功 notice from on_success")
        return super(MyTask, self).on_success(retval, task_id, args, kwargs)

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        r = redis.Redis(host=REDIS_ADDR, password=REDIS_PD, port=REDIS_PORT, db=ansible_result_redis_db)
        a = redis.Redis(host=REDIS_ADDR, password=REDIS_PD, port=REDIS_PORT, db=4)
        tid = args[0]
        rlist = r.lrange(tid, 0, -1)
        try:
            at = AnsibleTasks.objects.filter(AnsibleID=tid)[0]
            at.AnsibleResult = json.dumps([json.loads(i.decode()) for i in rlist])
            ct = a.get('celery-task-meta-%s' % at.CeleryID).decode()
            at.CeleryResult = ct
            at.save()
        except:
            pass
        return super(MyTask, self).on_failure(exc, task_id, args, kwargs, einfo)


@appCelery.task()
def ansibleExec(tid, groupname, tasks=[], extra_vars={}):
    try:
        AnsibleApi(tid, groupname, tasks, sources, extra_vars)
        return 'success'
    except Exception as e:
        celery_logger.error(e)
        return 'fail'


@appCelery.task(bind=True, base=MyTask)
def ansiblePlayBook(self, tid, pb, extra_vars={}, **kw):
    psources = kw.get('sources') or extra_vars.get('sources') or sources
    try:
        AnsiblePlaybookApi(tid, ["../ansible_files/playbooks/%s" % pb], psources, extra_vars)
        return 'success'
    except Exception as e:
        celery_logger.error(e)
        return 'fail'


# 执行结束，结果保持至db
@appCelery.task(bind=True)  # bind=True 用于将自身信息传入函数
def syncAnsibleResult(self, ret, *a, **kw):
    c = AsyncResult(self.request.get('parent_id'))  # 获取父任务 id
    celery_logger.info(c.result)
    tid = kw.get('tid', None)
    if tid:
        r = redis.Redis(host=REDIS_ADDR, password=REDIS_PD, port=REDIS_PORT, db=ansible_result_redis_db)
        a = redis.Redis(host=REDIS_ADDR, password=REDIS_PD, port=REDIS_PORT, db=4)
        rlist = r.lrange(tid, 0, -1)  # 获取 ansible 执行结果
        at = AnsibleTasks.objects.filter(AnsibleID=tid)[0]  # 匹配到对于的任务对象
        at.AnsibleResult = json.dumps([json.loads(i.decode()) for i in rlist])
        ct = a.get('celery-task-meta-%s' % at.CeleryID).decode()
        at.CeleryResult = ct
        at.save()
    else:
        pass


@appCelery.task()  #
def test_celery(msg):
    return msg


if __name__ == "__main__":
    appCelery.worker_main()
